package com.agileboot.medicine.controller.algorithm;

import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.Fuzziness;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilder;
import org.elasticsearch.index.query.functionscore.ScoreFunctionBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;

import java.util.Collections;

public class SimilaritySearch {
    @Autowired
    private RestHighLevelClient restHighLevelClient;

    public static void main(String[] args) throws Exception {
        // 创建 Elasticsearch 高级客户端
//        RestHighLevelClient client = new RestHighLevelClient("");

        // 要比较的文本
        String queryText = "your_text_to_compare";

        // 构建搜索请求
        SearchRequest searchRequest = new SearchRequest("your_index");
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();

        // 设置查询条件，使用 Function Score Query
        searchSourceBuilder.query(QueryBuilders.functionScoreQuery(
                QueryBuilders.matchQuery("content", queryText)
                        .fuzziness(Fuzziness.AUTO),  // 允许一些模糊匹配
                (ScoreFunctionBuilder) Collections.singletonList(
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(
                                ScoreFunctionBuilders.scriptFunction(
                                        new Script(
                                                ScriptType.INLINE,
                                                "painless",
                                                "cosineSimilarity(params.query_vector, 'content_vector') + 1.0",
                                                Collections.singletonMap("query_vector", queryText)
                                        )
                                )
                        )
                )
        ));

        // 设置排序规则，按相似度降序排列
        searchSourceBuilder.sort(SortBuilders.scoreSort().order(SortOrder.DESC));

        // 设置返回结果数量，这里可以根据实际情况调整
        searchSourceBuilder.size(1);

        // 执行搜索
        searchRequest.source(searchSourceBuilder);
//        SearchResponse searchResponse = restHighLevelClient.search(searchRequest, RequestOptions.DEFAULT);

        // 处理搜索结果
        // 获取相似度最高的文档
//        if (searchResponse.getHits().getTotalHits().value > 0) {
//            String mostSimilarDocumentId = searchResponse.getHits().getAt(0).getId();
//            System.out.println("Most similar document ID: " + mostSimilarDocumentId);
//        } else {
//            System.out.println("No similar documents found.");
//        }

        // 关闭 Elasticsearch 高级客户端
//        restHighLevelClient.close();
    }
}
